# coding: utf-8
# implementation of artificial bee colony (ABC) for TSP
import os
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import sample, shuffle, uniform
import argparse
from time import time, gmtime, strftime, localtime

from TSP import TSP

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)

parser = argparse.ArgumentParser()
parser.add_argument("--case", required=True, type=str, help="The case name!")
args = parser.parse_args()

case = args.case

instance = TSP(path=f'../../data/TSPLIB/{case}.tsp.txt')
Distance = instance.distances  # distance matrix
V = deepcopy(instance.cities)  # city set


def make_dir(path):
    path = path.strip().rstrip("\\")
    if os.path.exists(path):
        print(path + ' already exists.')
        return False
    else:
        os.makedirs(path)
        return True


class Solution:
    def __init__(self):
        self.route = None
        self.length = None
        self.stagnation = 0
        # initialize a solution by minimal insertion
        unvisited = deepcopy(V)
        d = sample(unvisited, k=1)[0]
        self.route = [d, d]
        self.length = 0
        unvisited.remove(d)
        shuffle(unvisited)
        for city in unvisited:
            self.optimal_insert(city)

    def calculate_length(self):
        self.length = 0
        for i in range(1, len(self.route)):
            self.length += Distance[self.route[i - 1]][self.route[i]]
        self.length += Distance[self.route[-1]][self.route[0]]  # add the depot to form the circle

    def optimal_insert(self, city):
        assert self.route is not None and self.length is not None
        opt_i, opt_cost = -1, 1e99
        for i in range(1, len(self.route)):
            cost = Distance[self.route[i - 1]][city] + Distance[city][self.route[i]] \
                   - Distance[self.route[i - 1]][self.route[i]]
            if opt_cost > cost:
                opt_i, opt_cost = i, cost
        self.route.insert(opt_i, city)
        self.length += opt_cost

    def make_neighbor(self, p):
        neighbor = deepcopy(self)
        neighbor.route = [self.route[0]]
        removal = []
        for city in self.route[1:-1]:  # remove cities at prob
            if uniform(0, 1) < p:
                neighbor.route.append(city)
            else:
                removal.append(city)
        neighbor.route.append(self.route[-1])
        shuffle(removal)
        for city in removal:  # re-insert
            neighbor.optimal_insert(city)
        neighbor.calculate_length()
        return neighbor


METHOD = 'ABC for TSP'
plt.figure(figsize=(10, 4))
max_epoch = 5
max_gen = 10000
NP = 50  # size of population
Pc = 0.8  # remove probability
if __name__ == '__main__':
    date = strftime("%Y-%m-%d", localtime())
    save_path = f'./results/{date}/{case}'
    make_dir(save_path)
    start_time = time()
    fitness = []
    for epoch in range(1, max_epoch + 1):
        POP = []
        # Initialization.
        for _ in range(NP):
            POP.append(Solution())
        pBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest_list = []
        pBest_list = []
        # Progress
        for gen in range(1, max_gen + 1):
            # Employed bee search.
            for _ in range(NP):
                neighbor = POP[_].make_neighbor(p=Pc)  # generate neighbor
                if neighbor.length < POP[_].length:
                    POP[_] = deepcopy(neighbor)
                else:
                    POP[_].stagnation += 1
            # Onlooker bee search.
            for _ in range(NP):
                a, b = sample(list(range(0, NP)), 2)
                if POP[a].length > POP[b].length:
                    a, b = b, a  # tournament selection
                neighbor = POP[a].make_neighbor(p=Pc)  # generate neighbor out of a
                if neighbor.length < POP[_].length:
                    POP[_] = deepcopy(neighbor)
                else:
                    POP[_].stagnation += 1
            # Scout bee watch.
            for _ in range(NP):
                if POP[_].stagnation > NP:  # stagnation limit reached
                    POP[_] = deepcopy(Solution())
            # Record gBest
            pBest = deepcopy(min(POP, key=lambda _: _.length))
            if gBest.length > pBest.length:
                gBest = deepcopy(pBest)
            pBest_list.append(pBest.length)
            gBest_list.append(gBest.length)
            print(f'method={METHOD}, case={case}, NP={NP}, Pc={Pc}. '
                  f'At epoch {epoch} gen {gen}: {round(gBest.length, 2)}', end='\r')
            # plot curve and route
            plt.clf()
            plt.subplot(1, 2, 1)
            plt.plot(range(len(gBest_list)), gBest_list, label='gBest')
            plt.legend()
            plt.title(f'length: {gBest.length:.2f}')
            plt.subplot(1, 2, 2)
            instance.draw(plt, route=gBest.route)
            plt.pause(0.001)
        plt.show()

        fitness.append(gBest_list[-1])
        mode = 'w' if epoch == 1 else 'ab'
        # ---------> Change save file here! <---------
        with open(save_path + rf'/{METHOD} for TSP.txt', mode) as f:
            np.savetxt(f, np.array(gBest_list).reshape(1, len(gBest_list)))

    print(f'method={METHOD}, case={case}, NP={NP}, Pc={Pc}. '
          f'At epoch {max_epoch} gen {max_gen}: {round(np.array(fitness).mean(), 2)}')
    t = strftime('%H:%M:%S', gmtime(time() - start_time))
    print(f'Time cost {t}')
